YOLO后处理

首发于个人博客

理论分析

system.png

YOLO从v2版本开始重新启用anchor box,YOLOv2网络的网络输出为尺寸为[b,125,13,13]的tensor,要将这个Tensor变为最终的输出结果,还需要以下的处理:

  • 解码:从Tensor中解析出所有框的位置信息和类别信息
  • NMS:筛选最能表现物品的识别框

解码过程

解码之前,需要明确的是每个候选框需要5+class_num个数据,分别是相对位置x,y,相对宽度w,h,置信度c和class_num个分类结果,YOLOv2-voc中class_num=20,即每个格点对应5个候选框,每个候选框有5+20=25个参数,这就是为什么输出Tensor的最后一维为5*(20+5)=125。

tensor.png

上图为一个框所需要的所有数据构成,假设这个框是位于格点X,Y的,对应的anchor box大小为W,H,位置相关参数的处理方法如下所示,其中,T_x,T_y分别是输出Tensor在长宽上的值,这里T_x = 13,T_y = 13P_x,P_y分别为原图片的长和宽:
x_r = \cfrac{sigmoid(x) + X}{T_x} \times P_x\\ y_r = \cfrac{sigmoid(x) + Y}{T_y} \times P_y\\ w_r = e^{w} \times W \\ h_r = e^{h} \times H
置信度和类别信息处理方法如下所示:
c_{r} = sigmoid(c) \times max\{softmax(class)\} \\ class\_id = argmax(class)
当格点置信度大于某个阈值时,认为该格点有物体,物体类别为class_id对应的类别

NMS

NMS为非最大值抑制,用在YOLO系统中的含义指从多个候选框标记同一个物品时,从中选择最合适的候选框。其基本思维很简单:使用置信度最高的候选框标记一个物体,若其他候选框与该候选框的IOU超过一个阈值,则认为其他候选框与该候选框标记的是同一个物体,丢弃其他候选框。

具体实现时,可以将所有候选框进行排序,置信度高的在前,置信度低的在后。从置信度高的候选框开始遍历所有候选框,对于某一个候选框,将之后所有的候选框与其计算IOU,若IOU高于一个阈值,则丢弃置信度低的候选框。算法流程图如下所示:

nms.png

代码分析

这里选择的是marvis开源的基于Pytorch的YOLOv2代码,其优势在于所有的部分均使用Python实现,没有使用Cython,无需编译即可使用,且依赖较少,文件管理比较扁平。

解码部分

解码部分在utils.py文件中,由get_region_boxes函数实现。首先是准备部分,这里首先获取了输出的相关信息,yolo-voc网络下有b为batch,预测模式下一般为1,h=w=13。随后reshape了输出,其维度变为(25,13*13*5),改变维度的目的是方便后面处理的索引。

def get_region_boxes(output, conf_thresh, num_classes, anchors, num_anchors, only_objectness=1, validation=False):    
    anchor_step = len(anchors) / num_anchors
    if output.dim() == 3:
        output = output.unsqueeze(0)
    batch = output.size(0)
    assert(output.size(1) == (5 + num_classes) * num_anchors)
    h = output.size(2)
    w = output.size(3)
    output = output.view(batch * num_anchors, 5 + num_classes, h * w).transpose(
        0, 1).contiguous().view(5 + num_classes, batch * num_anchors * h * w)
    all_boxes = []

随后是处理x,y的部分,xs和ys就是处理后的候选框中心点相对坐标,grid_x和grid_y与output[0]shape相同,分别表示对应output位置的候选框所属的格点坐标X与Y,这里的xs和ys实现了上述公式中的xs = sigmoid(x) + Xys = sigmoid(y) + Y

    grid_x = torch.linspace(0, w - 1, w).repeat(h, 1).repeat(batch *
                                                             num_anchors, 1, 1).view(batch * num_anchors * h * w).cuda()
    grid_y = torch.linspace(0, h - 1, h).repeat(w, 1).t().repeat(
        batch * num_anchors, 1, 1).view(batch * num_anchors * h * w).cuda()
    print("outputs shape", output.shape)
    xs = torch.sigmoid(output[0]) + grid_x
    ys = torch.sigmoid(output[1]) + grid_y

之后为处理w,h的部分,与处理x,y的部分类似,最终ws和hs为修正后的物品尺寸信息,实现了ws = e^w\cdot Whs = e^{h} \cdot H。其中W和H分别为当前anchor box的建议尺寸。

    anchor_w = torch.Tensor(anchors).view(
        num_anchors, anchor_step).index_select(1, torch.LongTensor([0]))
    anchor_h = torch.Tensor(anchors).view(
        num_anchors, anchor_step).index_select(1, torch.LongTensor([1]))
    anchor_w = anchor_w.repeat(batch, 1).repeat(
        1, 1, h * w).view(batch * num_anchors * h * w).cuda()
    anchor_h = anchor_h.repeat(batch, 1).repeat(
        1, 1, h * w).view(batch * num_anchors * h * w).cuda()
    ws = torch.exp(output[2]) * anchor_w
    hs = torch.exp(output[3]) * anchor_h

接下来是获取置信度的部分和类别部分,获取该anchor box的置信度为det_confs=sigmoid(c)。随后处理类别信息,先对类别信息对应的数据做softmax操作,随后获取其最大值cls_max_confs和最大值所在的位置cls_max_ids,其中位置cls_max_ids对应每个anchor box框住的“物品”的类别。

    det_confs = torch.sigmoid(output[4])
    cls_confs = torch.nn.Softmax()(
        Variable(output[5:5 + num_classes].transpose(0, 1))).data
    cls_max_confs, cls_max_ids = torch.max(cls_confs, 1)
    cls_max_confs = cls_max_confs.view(-1)
    cls_max_ids = cls_max_ids.view(-1)

随后是一些其他的处理过程,例如获取格点数量sz_hw,anchor box的数量sz_hwa等,函数convert2cpu是在CPU上复制一个该数据,注意这里是拷贝,并不是将数据从GPU转移到CPU上。

    sz_hw = h * w
    sz_hwa = sz_hw * num_anchors
    det_confs = convert2cpu(det_confs)
    cls_max_confs = convert2cpu(cls_max_confs)
    cls_max_ids = convert2cpu_long(cls_max_ids)
    xs = convert2cpu(xs)
    ys = convert2cpu(ys)
    ws = convert2cpu(ws)
    hs = convert2cpu(hs)
    if validation:
        cls_confs = convert2cpu(cls_confs.view(-1, num_classes))

随后是一个解码的大循环,分析见下面的注释

    for b in range(batch):
        boxes = []
        # boxes为容纳所有候选框的list
        for cy in range(h):
            for cx in range(w):
                for i in range(num_anchors):
                    # 遍历每一个anchor box,这里访问位于格点cx,cy的第i个anchor box
                    ind = b * sz_hwa + i * sz_hw + cy * w + cx
                    # 获取该anchor box在det_conf中对应的index
                    det_conf = det_confs[ind]
                    if only_objectness:
                        conf = det_confs[ind]
                    else:
                        conf = det_confs[ind] * cls_max_confs[ind]
                    # 处理置信度

                    if conf > conf_thresh:
                        # 若置信度大于阈值,则认为该anchor box有效
                        bcx = xs[ind]
                        bcy = ys[ind]
                        bw = ws[ind]
                        bh = hs[ind]
                        cls_max_conf = cls_max_confs[ind]
                        cls_max_id = cls_max_ids[ind]
                        # 获取所有相关信息,包括长,宽,位置,置信度和类别
                        box = [bcx / w, bcy / h, bw / w, bh / h,
                               det_conf, cls_max_conf, cls_max_id]
                        # 处理数据,其中位置信息x,y,尺寸信息w,h均归一化,使其与输入图片尺寸解耦
                        if (not only_objectness) and validation:
                            for c in range(num_classes):
                                tmp_conf = cls_confs[ind][c]
                                if c != cls_max_id and det_confs[ind] * tmp_conf > conf_thresh:
                                    box.append(tmp_conf)
                                    box.append(c)
                        boxes.append(box)
                        # 将处理好的anchor box信息保存在boxes中
        all_boxes.append(boxes)
        return all_boxes

NMS部分

NMS也在utils.py中,函数名为nms。该函数中,首先实现对所有候选框的排序。这里使用det_confs获取了置信度从大到小的anchor box的坐标位置sortIds

def nms(boxes, nms_thresh):
    if len(boxes) == 0:
        return boxes

    det_confs = torch.zeros(len(boxes))
    for i in range(len(boxes)):
        det_confs[i] = 1 - boxes[i][4]

    _, sortIds = torch.sort(det_confs)

随后实现候选框的筛选,从高置信度的候选框开始遍历,对于每个候选框boxes[sortIds[i]],遍历所有置信度低于该候选框且置信度不为0(置信度为0表示该候选框被抛弃)的候选框,若低置信度候选框与高置信度候选框的IOU大于阈值,则抛弃低置信度候选框。

    out_boxes = []
    for i in range(len(boxes)):
        # 按置信度从高到低遍历
        box_i = boxes[sortIds[i]]
        if box_i[4] > 0:
            # 置信度大于0表示该候选框没有在之前的筛选中被抛弃
            out_boxes.append(box_i)
            for j in range(i + 1, len(boxes)):
                # 遍历所有置信度低于该候选框的候选框
                box_j = boxes[sortIds[j]]
                if bbox_iou(box_i, box_j, x1y1x2y2=False) > nms_thresh:
                    # 若置信度低的候选框与该候选框IOU大于一定值,抛弃低置信度候选框
                    box_j[4] = 0
    return out_boxes
©著作权归作者所有,转载或内容合作请联系作者
  • 序言:七十年代末,一起剥皮案震惊了整个滨河市,随后出现的几起案子,更是在滨河造成了极大的恐慌,老刑警刘岩,带你破解...
    沈念sama阅读 159,458评论 4 363
  • 序言:滨河连续发生了三起死亡事件,死亡现场离奇诡异,居然都是意外死亡,警方通过查阅死者的电脑和手机,发现死者居然都...
    沈念sama阅读 67,454评论 1 294
  • 文/潘晓璐 我一进店门,熙熙楼的掌柜王于贵愁眉苦脸地迎上来,“玉大人,你说我怎么就摊上这事。” “怎么了?”我有些...
    开封第一讲书人阅读 109,171评论 0 243
  • 文/不坏的土叔 我叫张陵,是天一观的道长。 经常有香客问我,道长,这世上最难降的妖魔是什么? 我笑而不...
    开封第一讲书人阅读 44,062评论 0 207
  • 正文 为了忘掉前任,我火速办了婚礼,结果婚礼上,老公的妹妹穿的比我还像新娘。我一直安慰自己,他们只是感情好,可当我...
    茶点故事阅读 52,440评论 3 287
  • 文/花漫 我一把揭开白布。 她就那样静静地躺着,像睡着了一般。 火红的嫁衣衬着肌肤如雪。 梳的纹丝不乱的头发上,一...
    开封第一讲书人阅读 40,661评论 1 219
  • 那天,我揣着相机与录音,去河边找鬼。 笑死,一个胖子当着我的面吹牛,可吹牛的内容都是我干的。 我是一名探鬼主播,决...
    沈念sama阅读 31,906评论 2 313
  • 文/苍兰香墨 我猛地睁开眼,长吁一口气:“原来是场噩梦啊……” “哼!你这毒妇竟也来了?” 一声冷哼从身侧响起,我...
    开封第一讲书人阅读 30,609评论 0 200
  • 序言:老挝万荣一对情侣失踪,失踪者是张志新(化名)和其女友刘颖,没想到半个月后,有当地人在树林里发现了一具尸体,经...
    沈念sama阅读 34,379评论 1 246
  • 正文 独居荒郊野岭守林人离奇死亡,尸身上长有42处带血的脓包…… 初始之章·张勋 以下内容为张勋视角 年9月15日...
    茶点故事阅读 30,600评论 2 246
  • 正文 我和宋清朗相恋三年,在试婚纱的时候发现自己被绿了。 大学时的朋友给我发了我未婚夫和他白月光在一起吃饭的照片。...
    茶点故事阅读 32,085评论 1 261
  • 序言:一个原本活蹦乱跳的男人离奇死亡,死状恐怖,灵堂内的尸体忽然破棺而出,到底是诈尸还是另有隐情,我是刑警宁泽,带...
    沈念sama阅读 28,409评论 2 254
  • 正文 年R本政府宣布,位于F岛的核电站,受9级特大地震影响,放射性物质发生泄漏。R本人自食恶果不足惜,却给世界环境...
    茶点故事阅读 33,072评论 3 237
  • 文/蒙蒙 一、第九天 我趴在偏房一处隐蔽的房顶上张望。 院中可真热闹,春花似锦、人声如沸。这庄子的主人今日做“春日...
    开封第一讲书人阅读 26,088评论 0 8
  • 文/苍兰香墨 我抬头看了看天上的太阳。三九已至,却和暖如春,着一层夹袄步出监牢的瞬间,已是汗流浃背。 一阵脚步声响...
    开封第一讲书人阅读 26,860评论 0 195
  • 我被黑心中介骗来泰国打工, 没想到刚下飞机就差点儿被人妖公主榨干…… 1. 我叫王不留,地道东北人。 一个月前我还...
    沈念sama阅读 35,704评论 2 276
  • 正文 我出身青楼,却偏偏与公主长得像,于是被迫代替她去往敌国和亲。 传闻我的和亲对象是个残疾皇子,可洞房花烛夜当晚...
    茶点故事阅读 35,608评论 2 270

推荐阅读更多精彩内容

  • 1.监听输入法的返回事件 关键是EditText的public boolean onKeyPreIme(int k...
    小山包阅读 476评论 0 0
  • 总有一些传奇的社群 若不走进 也许永远都不会接触 因为超级猩猩的缘起 开始接触到背后这个叫做莱美的组织 来自新西兰...
    RayofLight阅读 710评论 1 0